import matplotlib.pyplot as plt
import numpy as np

plt.rcParams['font.family'] = 'Times New Roman'

fig, ax = plt.subplots(figsize=(5, 5), dpi=300, constrained_layout=True)
fig.patch.set_facecolor('white')
ax.set_facecolor('white')

ax.set_xlim(-1.0, 1.0)
ax.set_ylim(-1.0, 1.0)
ax.margins(0)

ax.set_xticks([])
ax.set_yticks([])
ax.set_xlabel("")
ax.set_ylabel("")
ax.set_anchor('C')

for spine in ax.spines.values():
    spine.set_visible(False)

ax.plot([-1.0, 0], [0, 0], color='black', linewidth=1.0)
ax.plot([0, 0], [-1.0, 0], color='black', linewidth=1.0)
ax.annotate("", xy=(0.98, 0), xytext=(0, 0), xycoords='data',
            arrowprops=dict(arrowstyle='-|>', lw=1.0, color='black', mutation_scale=12))
ax.annotate("", xy=(0, 0.98), xytext=(0, 0), xycoords='data',
            arrowprops=dict(arrowstyle='-|>', lw=1.0, color='black', mutation_scale=12))
ax.text(0.98, 0.05, 'Valence', fontsize=12, fontweight='bold', ha='right', va='bottom')
ax.text(-0.05, 0.98, 'Arousal', fontsize=12, fontweight='bold', ha='right', va='top')

states = {
    'Interested': (0.5, 0.5),
    'Calmly Focused': (0.5, -0.5),
    'Restless': (-0.5, 0.5),
    'Bored': (-0.5, -0.5),
}

palette = {
    'Interested': '#0072B2',       # blue
    'Calmly Focused': '#009E73',   # green
    'Restless': '#D55E00',         # vermillion
    'Bored': '#E69F00',            # orange
}

for state, (x, y) in states.items():
    c = palette.get(state, 'black')
    ax.plot(x, y, marker='o', markersize=4.5, color=c, zorder=3)
    dx = 0.035 if x >= 0 else -0.035
    dy = 0.035 if y >= 0 else -0.035
    ha = 'left' if dx > 0 else 'right'
    va = 'bottom' if dy > 0 else 'top'
    ax.text(x + dx, y + dy, state,
            fontsize=12, fontweight='semibold', ha=ha, va=va, color=c)

ax.set_aspect('equal')

plt.savefig('student_states_valence_arousal.png', dpi=600, facecolor='white')
plt.savefig('student_states_valence_arousal.svg', facecolor='white')